import pickle
import csv
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '1'

import numpy as np
import tensorflow as tf


@tf.function(jit_compile=False)
def run_no_xla(x):
    return tf.math.sign(x)


@tf.function(jit_compile=True)
def run_xla(x):
    return tf.math.sign(x)


a = np.array([-3.4e+307 + 0.j, -6.32e+307 + 0.j], dtype=np.complex128)
result1 = run_no_xla(a)
print(result1)
result2 = run_xla(a)
print(result2)
